# face_recognizer.py - 人脸识别模块（提高识别准确率）
import cv2
import numpy as np
from config import get_user_names

MODEL_PATH = "trainer/trainer.yml"
CASCADE_PATH = "HaarCascade/haarcascade_frontalface_default.xml"
DATA_PATH = "data/"
MIN_FACE_RATIO = 0.1
CONFIDENCE_THRESHOLD = 90

class FaceRecognizer:
    def __init__(self):
        # 初始化模型和分类器
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        try:
            self.recognizer.read(MODEL_PATH)
        except:
            print("模型文件不存在或格式错误")
            
        self.face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
        
        # 添加识别历史记录，用于提高准确率
        self.recognition_history = []
        self.history_size = 5  # 保存最近5次识别结果
        
        print("人脸识别模块初始化完成")
    
    def process_frame(self, frame):
        """处理单个视频帧并返回识别结果"""
        # 转换为灰度图
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # 计算最小人脸尺寸（基于当前视频分辨率）
        height, width = frame.shape[:2]
        min_face_size = (
            int(width * MIN_FACE_RATIO),
            int(height * MIN_FACE_RATIO)
        )

        # 人脸检测
        faces = self.face_cascade.detectMultiScale(
            gray,
            scaleFactor=1.2,
            minNeighbors=5,
            minSize=min_face_size
        )
        
        face_names = []
        face_locations = []
        confidence_values = []
        
        # 获取最新的用户列表
        user_names = get_user_names()
        
        for (x, y, w, h) in faces:
            # 人脸识别
            id_num, confidence = self.recognizer.predict(gray[y:y + h, x:x + w])
            
            # 使用历史记录提高识别准确率
            if confidence < CONFIDENCE_THRESHOLD and id_num < len(user_names):
                name = user_names[id_num]
                
                # 添加到历史记录
                self.recognition_history.append(name)
                if len(self.recognition_history) > self.history_size:
                    self.recognition_history.pop(0)
                
                # 使用历史记录中最频繁出现的名称
                if len(self.recognition_history) == self.history_size:
                    from collections import Counter
                    most_common = Counter(self.recognition_history).most_common(1)
                    if most_common:
                        name = most_common[0][0]
            else:
                name = "unknown"
                # 重置历史记录
                self.recognition_history = []
            
            face_names.append(name)
            face_locations.append((x, y, w, h))
            confidence_values.append(round(100 - confidence))
            
            # 绘制人脸矩形框
            color = (0, 0, 255) if name == "unknown" else (0, 255, 0)
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            
            # 绘制识别结果
            cv2.putText(frame, name, (x + 5, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
            cv2.putText(frame, f"{confidence_values[-1]}%", (x + 5, y + h - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
        
        return frame, face_locations, face_names, confidence_values